package com.guozinhs.auto.match;

import cn.hutool.core.util.ObjectUtil;
import com.guozinhs.auto.util.image.ImageUtils;

import java.awt.*;
import java.awt.image.BufferedImage;
import java.util.List;
import java.util.Queue;
import java.util.*;
import java.util.stream.Collectors;

/**
 * 图像模板匹配
 *
 * @Author: GuoZi
 * @Date: 2022/4/2 16:00
 */
public class MyImageMatcher {

    private static final int STEP = 3;

    /**
     * 图像模板匹配
     *
     * @param resourceImage 原图像
     * @param templateImage 模板图像
     * @return 最匹配的位置信息
     */
    public static PointInfo match(BufferedImage resourceImage, BufferedImage templateImage) {
        int floor = 2;
        int scale = (int) Math.pow(STEP, floor);
        if (Objects.isNull(resourceImage) || Objects.isNull(templateImage)) {
            return null;
        }
        int width = templateImage.getWidth();
        int height = templateImage.getHeight();
        if (width < (scale << 1) || height < (scale << 1)) {
            floor = 1;
        }
        int[][] resourceArr = parseGrayArr(resourceImage);
        int[][] templateArr = parseGrayArr(templateImage);
        List<PointInfo> points = match(resourceArr, templateArr, floor, 1);
        PointInfo pointInfo = points.get(0);
        Point point = pointInfo.getPoint();
        pointInfo.setMatchPercent(matchPercent(resourceImage, templateImage, point.x, point.y));
        return pointInfo;
    }

    /**
     * 绘制匹配结果(矩形边框)
     *
     * @param source        原图
     * @param template      匹配模板
     * @param matchPoint    匹配点
     * @param resultImgPath 保存的图片路径
     */
    public static void drawMatch(BufferedImage source, BufferedImage template, Point matchPoint, String resultImgPath) {
        // 在原图的匹配范围绘制绿色矩形
        Graphics graphics = source.getGraphics();
        graphics.setColor(new Color(0, 255, 0));
        graphics.drawRect(matchPoint.x, matchPoint.y, template.getWidth() - 1, template.getHeight() - 1);
        // 保存匹配日志
        ImageUtils.write(resultImgPath, source);
    }

    /**
     * 获取颜色的RGB值
     *
     * @param color 颜色
     * @return RGB值
     */
    private static int[] getRGB(Color color) {
        return new int[]{color.getRed(), color.getGreen(), color.getBlue()};
    }

    /**
     * 灰度值匹配, 金字塔搜索法
     *
     * @param resourceArr 原图像灰度值
     * @param templateArr 模板图像灰度值
     * @param floor       金字塔层数
     * @param current     当前匹配层数
     * @return 若干个最匹配的位置信息
     */
    private static List<PointInfo> match(int[][] resourceArr, int[][] templateArr, int floor, int current) {
        if (floor < 1) {
            return Collections.emptyList();
        }
        int size = (current == 1) ? 1 : 64;
        if (floor == current) {
            return match(resourceArr, templateArr, size);
        }
        List<PointInfo> pointInfoList = match(getSubArr(resourceArr), getSubArr(templateArr), floor, current + 1);
        return match(resourceArr, templateArr, size, pointInfoList);
    }

    /**
     * 灰度值转换
     *
     * @param image 图像对象
     * @return 二维数组
     */
    private static int[][] parseGrayArr(BufferedImage image) {
        int width = image.getWidth();
        int height = image.getHeight();
        int[][] result = new int[width][height];
        for (int x = 0; x < width; x++) {
            for (int y = 0; y < height; y++) {
                int rgb = image.getRGB(x, y);
                result[x][y] = getGray(rgb);
            }
        }
        return result;
    }

    /**
     * 获取子数组
     *
     * @param grayArr 灰度值二维数组
     * @return 子数组
     */
    private static int[][] getSubArr(int[][] grayArr) {
        int width = grayArr.length;
        int height = grayArr[0].length;
        int subWidth = width / STEP;
        int subHeight = height / STEP;
        int[][] result = new int[subWidth][subHeight];
        for (int x = 0; x < subWidth; x++) {
            for (int y = 0; y < subHeight; y++) {
                result[x][y] = getAverage(grayArr, x * STEP, y * STEP);
            }
        }
        return result;
    }

    /**
     * 获取平均灰度值
     *
     * @param grayArr 灰度值二维数组
     * @param offsetX 计算位置X
     * @param offsetY 计算位置Y
     * @return 平均值
     */
    private static int getAverage(int[][] grayArr, int offsetX, int offsetY) {
        int result = 0;
        int width = grayArr.length;
        int height = grayArr[0].length;
        for (int x = 0; x < STEP; x++) {
            for (int y = 0; y < STEP; y++) {
                if (offsetX + x >= width || offsetY + y >= height) {
                    result += 255;
                } else {
                    result += grayArr[offsetX + x][offsetY + y];
                }
            }
        }
        return result / (STEP * STEP);
    }

    /**
     * 灰度值匹配
     *
     * @param resourceArr 原灰度值数组
     * @param templateArr 模板灰度值数组
     * @param size        需要的位置信息数量
     * @return 指定数量的位置信息
     */
    private static List<PointInfo> match(int[][] resourceArr, int[][] templateArr, int size) {
        int tempWidth = templateArr.length;
        int tempHeight = templateArr[0].length;
        int widthRange = resourceArr.length - tempWidth;
        int heightRange = resourceArr[0].length - tempHeight;
        if (widthRange < 0) {
            throw new RuntimeException("模板宽度过长");
        }
        if (heightRange < 0) {
            throw new RuntimeException("模板高度过长");
        }
        PriorityQueue<PointInfo> queue = new PriorityQueue<>(size + 1);
        for (int x = 0; x <= widthRange; x++) {
            for (int y = 0; y <= heightRange; y++) {
                double matchPercent = matchPercent(resourceArr, templateArr, x, y);
                PointInfo info = new PointInfo();
                info.setPoint(new Point(x, y));
                info.setMatchPercent(matchPercent);
                queue.add(info);
                if (queue.size() > size) {
                    queue.poll();
                }
            }
        }
        return toList(queue);
    }


    /**
     * 灰度值匹配
     *
     * @param resourceArr   原灰度值数组
     * @param templateArr   模板灰度值数组
     * @param size          需要的位置信息数量
     * @param pointInfoList 指定的匹配位置信息
     * @return 指定数量的位置信息
     */
    private static List<PointInfo> match(int[][] resourceArr, int[][] templateArr, int size, List<PointInfo> pointInfoList) {
        if (ObjectUtil.isEmpty(pointInfoList)) {
            return Collections.emptyList();
        }
        List<Point> pointList = pointInfoList.stream().map(PointInfo::getPoint).filter(Objects::nonNull).collect(Collectors.toList());
        int tempWidth = templateArr.length;
        int tempHeight = templateArr[0].length;
        int widthRange = resourceArr.length - tempWidth;
        int heightRange = resourceArr[0].length - tempHeight;
        if (widthRange < 0) {
            throw new RuntimeException("模板宽度过长");
        }
        if (heightRange < 0) {
            throw new RuntimeException("模板高度过长");
        }
        Set<String> handledPoint = new HashSet<>();
        PriorityQueue<PointInfo> queue = new PriorityQueue<>(size + 1);
        // 精匹配
        for (Point point : pointList) {
            for (int x = -STEP; x < STEP; x++) {
                for (int y = -STEP; y < STEP; y++) {
                    int offsetX = point.x * STEP + x;
                    int offsetY = point.y * STEP + y;
                    if (offsetX < 0 || offsetY < 0 || offsetX > widthRange || offsetY > heightRange) {
                        // 超出边界
                        continue;
                    }
                    if (!handledPoint.add(String.valueOf(offsetX) + offsetY)) {
                        // 避免重复匹配
                        continue;
                    }
                    double matchPercent = matchPercent(resourceArr, templateArr, offsetX, offsetY);
                    PointInfo info = new PointInfo();
                    info.setPoint(new Point(offsetX, offsetY));
                    info.setMatchPercent(matchPercent);
                    queue.add(info);
                    if (queue.size() > size) {
                        queue.poll();
                    }
                }
            }
        }
        return toList(queue);
    }


    /**
     * 图像匹配
     * <p>
     * 按灰度值进行匹配, 并返回匹配的像素点数量比例
     *
     * @param source   原灰度值数组
     * @param template 模板灰度值数组
     * @param offsetX  x偏移
     * @param offsetY  y偏移
     * @return 匹配比例
     */
    private static double matchPercent(int[][] source, int[][] template, int offsetX, int offsetY) {
        int width = template.length;
        int height = template[0].length;
        int count = 0;
        int total = 0;
        for (int col = 0; col < width; col++) {
            for (int row = 0; row < height; row++) {
                total++;
                int s = source[col + offsetX][row + offsetY];
                int t = template[col][row];
                if (matchGray(s, t)) {
                    count++;
                }
            }
        }
        return count * 100.0 / total;
    }

    /**
     * 图像匹配
     * <p>
     * 按像素值进行匹配, 并返回匹配的像素点数量比例
     *
     * @param source   原图像
     * @param template 模板图像
     * @param offsetX  x偏移
     * @param offsetY  y偏移
     * @return 匹配比例
     */
    private static double matchPercent(BufferedImage source, BufferedImage template, int offsetX, int offsetY) {
        int width = template.getWidth();
        int height = template.getHeight();
        int count = 0;
        int total = width * height;
        for (int x = 0; x < width; x++) {
            for (int y = 0; y < height; y++) {
                boolean match = true;
                int[] s = getRGB(new Color(source.getRGB(offsetX + x, offsetY + y)));
                int[] t = getRGB(new Color(template.getRGB(x, y)));
                for (int i = 0; i < s.length; i++) {
                    double v = Math.abs(s[i] - t[i]);
                    if (v >= 32) {
                        match = false;
                        break;
                    }
                }
                if (match) {
                    count++;
                }
            }
        }
        return count * 100.0 / total;
    }

    /**
     * 灰度值匹配
     *
     * @param x 灰度值
     * @param y 另一个灰度值
     * @return 匹配结果
     */
    private static boolean matchGray(int x, int y) {
        return Math.abs(x - y) <= 32;
    }

    /**
     * RGB值转灰度值
     *
     * @param rgb RGB值
     * @return 灰度值
     */
    private static int getGray(int rgb) {
        int r = (rgb >> 16) & 0xFF;
        int g = (rgb >> 8) & 0xFF;
        int b = rgb & 0xFF;
        return (r + g + b) / 3;
    }

    /**
     * 队列转List集合
     *
     * @param queue 队列
     * @param <T>   数据类型
     * @return List集合
     */
    private static <T> List<T> toList(Queue<T> queue) {
        List<T> result = new ArrayList<>();
        while (!queue.isEmpty()) {
            result.add(queue.poll());
        }
        return result;
    }

}
